package com.aikg.kgrag_java.pojo.dto.neo4j;

import com.aikg.kgrag_java.pojo.entity.neo4j.Node;
import lombok.Data;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @ClassName: ForMeRetrievalDTO
 * @Description: 保存neo4j检索到的子图（节点及其之间的关系）
 * @Author: 李宗雨
 * @Date: 2025/1/5 14:54
 * @Version: 1.0
 **/
@Data
public class Neo4jRetrievalDTO {

    // 节点及其节点信息
    private List<Node> nodeInfoList = new ArrayList<>();

    // 节点间关系的三元组
    private List<Map<String, String>> relationTripletList = new ArrayList<>();


    public String nodeInfoListToString4Me(){
        StringBuilder temp = new StringBuilder();

        for(Node node: this.nodeInfoList){
            temp.append(node.toString4Me()).append(",");
        }
        if(!temp.isEmpty()) {
            temp.deleteCharAt(temp.length() - 1);
        }
        return temp.toString();
    }

    public String relationTripletListToString4Me() {
        StringBuilder temp = new StringBuilder();
        for(Map<String, String> triplet: this.relationTripletList){
            temp.append("{\"头实体\":").append("\"").append(triplet.get("头实体")).append("\",");
            temp.append("\"尾实体\":").append("\"").append(triplet.get("尾实体")).append("\",");
            temp.append("\"关系\":").append("\"").append(triplet.get("关系")).append("\"}");
        }
        return temp.toString();
    }

    public List<Map<String, String>> relationTripletListToFront(){
        // 使用Stream API进行转换
        List<Map<String, String>> transformedList = this.relationTripletList.stream()
                .map(map -> {
                    Map<String, String> newMap = new HashMap<>();
                    newMap.put("source", map.get("头实体"));
                    newMap.put("target", map.get("尾实体"));
                    newMap.put("value", map.get("关系"));
                    return newMap;
                })
                .collect(Collectors.toList());
        return transformedList;
    }
}
